package com.shigen.demo.design;

import cn.hutool.json.JSONUtil;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import lombok.AllArgsConstructor;
import lombok.Data;

/**
 * @author shigenfu
 * @date 2023/8/22 3:55 下午
 */
public class TreeNode {

    @Data
    @AllArgsConstructor
    static class Item {

        private Integer id;
        private String name;
        private Integer pid;
        private List<Item> children;

    }

    public static void main(String[] args) {

        ArrayList<Item> items = new ArrayList<Item>() {
            {
                add(new Item(111, "one", null, null));
                add(new Item(111-1, "one", null, null));
                add(new Item(222, "two", 111, null));
                add(new Item(333, "three", 111, null));
                add(new Item(444, "four", 333, null));
                add(new Item(444-1, "four", 333, null));
                add(new Item(444-2, "four", 444-1, null));
            }
        };

        List<Item> tree = getTree(items);
        String s = JSONUtil.parse(tree).toStringPretty();
        System.out.println(s);
    }

    private static List<Item> getTree(ArrayList<Item> items) {
        // 获得非顶级节点的数据并按照pid分组
        Map<Integer, List<Item>> nodeMap = items.stream().filter(item -> item.getPid() != null).collect(Collectors.groupingBy(Item::getPid));

        // 循环设置子节点
        items.forEach(item -> item.setChildren(nodeMap.get(item.id)));

        // 获得根节点的数据
        List<Item> treeNode = items.stream().filter(node -> node.getPid() == null).collect(Collectors.toList());
        return treeNode;
    }
}
